from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict


class State(TypedDict):
    input: str


def step_1(state):
    print("---Step 1---")
    return state


def step_2(state):
    print("---Step 2---")
    return state


def step_3(state):
    print("---Step 3---")
    return state


builder = StateGraph(State)
builder.add_node("step_1", step_1)
builder.add_node("step_2", step_2)
builder.add_node("step_3", step_3)
builder.add_edge(START, "step_1")
builder.add_edge("step_1", "step_2")
builder.add_edge("step_2", "step_3")
builder.add_edge("step_3", END)

# Set up memory
memory = MemorySaver()

# Add
graph = builder.compile(checkpointer=memory)

initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
for event in graph.stream(initial_input, thread, stream_mode="values"):
    print(event)

# try:
#     user_approval = input("Do you want to go to Step 3? (yes/no): ")
# except:
#     user_approval = "yes"
#
# if user_approval.lower() == "yes":
#     # If approved, continue the graph execution
#     for event in graph.stream(None, thread, stream_mode="values"):
#         print(event)
# else:
#     print("Operation cancelled by user.")